#include "ring_basis_extension.h"
#include <cmath>
#include <cstring>
#include <algorithm>

/*
 * @brief 构造函数，初始化基转换参数。
 *
 * 该函数初始化基转换参数，包括模数向量、基转换参数等。
 *
 * @param ringQ 目标环对象。
 * @param ringP 源环对象。
 */
FastBasisExtender::FastBasisExtender(Ring *ringQ, Ring *ringP)
{
    int MAXLK = 50;
    this->ringQ = ringQ;
    this->ringP = ringP;

    // 基转换参数初始化
    long K = ringP->Modulus.size();
    long L = ringQ->Modulus.size();
    if (L == 0 || L > MAXLK) {
        cout << " 间开辟错误 " << endl;
    }
    if (K == 0 || K > MAXLK) {
        cout << " 间开辟错误 " << endl;
    }
    qVec = new uint64_t[L]();
    pVec = new uint64_t[K]();
    for (long l = 0; l < L; l++) {
        qVec[l] = ringQ->Modulus[l];
    }
    for (long k = 0; k < K; k++) {
        pVec[k] = ringP->Modulus[k];
    }
    qHatModq = new uint64_t *[L]; // [l][i] (phat_i)_l mod p_i
    for (long l = 0; l < L; l++) {
        qHatModq[l] = new uint64_t[l + 1];
        for (long i = 0; i < l + 1; i++) {
            qHatModq[l][i] = 1;
            for (long j = 0; j < i; j++) {
                uint64_t temp = qVec[j] % qVec[i];
                MulMod(qHatModq[l][i], qHatModq[l][i], temp, qVec[i]);
            }
            for (long j = i + 1; j < l + 1; j++) {
                uint64_t temp = qVec[j] % qVec[i];
                MulMod(qHatModq[l][i], qHatModq[l][i], temp, qVec[i]);
            }
        }
    }

    qHatInvModq = new uint64_t *[L]; // [l][i] (phat_i)_l^-1 mod p_i
    for (long l = 0; l < L; l++) {
        qHatInvModq[l] = new uint64_t[l + 1];
        for (long i = 0; i < l + 1; i++) {
            qHatInvModq[l][i] = InvMod(qHatModq[l][i], qVec[i]);
        }
    }

    pHatModp = new uint64_t[K]; // [k] qhat_k mod q_k

    for (long k = 0; k < K; k++) {
        pHatModp[k] = 1;
        for (long j = 0; j < k; j++) {
            uint64_t temp = pVec[j] % pVec[k];
            MulMod(pHatModp[k], pHatModp[k], temp, pVec[k]);
        }
        for (long j = k + 1; j < K; j++) {
            uint64_t temp = pVec[j] % pVec[k];
            MulMod(pHatModp[k], pHatModp[k], temp, pVec[k]);
        }
    }

    pHatInvModp = new uint64_t[K]; // [k] qhat_k^-1 mod q_k
    pHatInvModp_shoup = new uint64_t[K];
    for (long k = 0; k < K; k++) {
        pHatInvModp[k] = InvMod(pHatModp[k], pVec[k]);
        pHatInvModp_shoup[k] = x_Shoup(pHatInvModp[k], pVec[k]);
    }

    qHatModp = new uint64_t **[L]; // [l] [i] [k]  (phat_i)_l mod q_k
    for (long l = 0; l < L; l++) {
        qHatModp[l] = new uint64_t *[l + 1];
        for (long i = 0; i < l + 1; i++) {
            qHatModp[l][i] = new uint64_t[K];
            for (long k = 0; k < K; k++) {
                qHatModp[l][i][k] = 1;
                for (long j = 0; j < i; j++) {
                    uint64_t temp = qVec[j] % pVec[k];
                    MulMod(qHatModp[l][i][k], qHatModp[l][i][k], temp, pVec[k]);
                }
                for (long j = i + 1; j < l + 1; j++) {
                    uint64_t temp = qVec[j] % pVec[k];
                    MulMod(qHatModp[l][i][k], qHatModp[l][i][k], temp, pVec[k]);
                }
            }
        }
    }

    pHatModq = new uint64_t *[K]; // [k][i] qhat_k mod p_i
    for (long k = 0; k < K; k++) {
        pHatModq[k] = new uint64_t[L];
        for (long i = 0; i < L; i++) {
            pHatModq[k][i] = 1;
            for (long s = 0; s < k; s++) {
                uint64_t temp = pVec[s] % qVec[i];
                MulMod(pHatModq[k][i], pHatModq[k][i], temp, qVec[i]);
            }
            for (long s = k + 1; s < K; s++) {
                uint64_t temp = pVec[s] % qVec[i];
                MulMod(pHatModq[k][i], pHatModq[k][i], temp, qVec[i]);
            }
        }
    }

    PModq = new uint64_t[L]; // [i] qprod mod p_i
    for (long i = 0; i < L; i++) {
        PModq[i] = 1;
        for (long k = 0; k < K; k++) {
            uint64_t temp = pVec[k] % qVec[i];
            MulMod(PModq[i], PModq[i], temp, qVec[i]);
        }
    }

    PInvModq = new uint64_t[L]; // [i] qprod^-1 mod p_i
    PInvModqShoup = new uint64_t[L];
    for (long i = 0; i < L; i++) {
        PInvModq[i] = InvMod(PModq[i], qVec[i]);
        PInvModqShoup[i] = x_Shoup(PInvModq[i], qVec[i]);
    }

    QModp = new uint64_t *[L];
    for (long i = 0; i < L; i++) {
        QModp[i] = new uint64_t[K];
        for (long k = 0; k < K; k++) {
            QModp[i][k] = 1;
            for (long j = 0; j < i + 1; j++) {
                uint64_t temp = qVec[j] % pVec[k];
                MulMod(QModp[i][k], QModp[i][k], temp, pVec[k]);
            }
        }
    }
    QInvModp = new uint64_t *[L];
    for (long i = 0; i < L; i++) {
        QInvModp[i] = new uint64_t[K];
        for (long k = 0; k < K; k++) {
            QInvModp[i][k] = InvMod(QModp[i][k], pVec[k]);
        }
    }

    qInvModq = new uint64_t *[L]; // [i][j] p_i^-1 mod p_j
    for (long i = 0; i < L; i++) {
        qInvModq[i] = new uint64_t[L];
        for (long j = 0; j < i; j++) {
            qInvModq[i][j] = InvMod(qVec[i], qVec[j]);
        }
        for (long j = i + 1; j < L; j++) {
            qInvModq[i][j] = InvMod(qVec[i], qVec[j]);
        }
    }

    baseConvPQ2Q = new uint64_t[K * ringQ->N]();
    polypoolQ = ringQ->NewPoly();
    BASECONVERT = (L > K ? L : K) + 1;
}

/*
 * @brief 将多项式 p1Q 的系数从 p0Q 中分解并分配到 p1Q 和 p1P 中。
 *
 * 该函数将 p0Q 中的系数分配到 p1Q 和 p1P 中，以实现基转换。
 *
 * @param levelQ 目标多项式的级数。
 * @param levelP 源多项式的级数。
 * @param alpha 分解的块大小。
 * @param beta 分解的块数量。
 * @param p0Q 输入多项式。
 * @param p1Q 输出多项式。
 * @param p1P 输出多项式。
 */
void Decomposer::DecomposeAndSplit(int levelQ, int levelP, int alpha, int beta, Poly *p0Q, Poly *p1Q, Poly *p1P)
{
    int lvlQStart = beta * alpha;
#pragma omp parallel for collapse(2)
    for (int j = 0; j < levelQ + 1; j++) {
        for (uint64_t k = 0; k < p1Q->Coeffs[0].size(); k++) {
            p1Q->Coeffs[j][k] = p0Q->Coeffs[lvlQStart][k];
        }
    }
#pragma omp parallel for collapse(2)
    for (int j = 0; j < levelP + 1; j++) {
        for (uint64_t k = 0; k < p1P->Coeffs[0].size(); k++) {
            p1P->Coeffs[j][k] = p0Q->Coeffs[lvlQStart][k];
        }
    }
}
/*
 * @brief 将基于 QP 的多项式转换为基于 Q 的多项式。
 *
 * 该函数实现了将两个多项式 `p1P` 和 `p1Q`（分别基于 P 和 Q）转换为一个基于 Q 的多项式 `p2Q`。
 *
 * @param levelQ Q 级数（转换级数）。
 * @param levelP P 级数（转换级数）。
 * @param p1Q 输入多项式 p1Q，基于 Q。
 * @param p1P 输入多项式 p1P，基于 P。
 * @param p2Q 输出多项式 p2Q，转换后的结果，基于 Q。
 */
void FastBasisExtender::BaseConvertQP2Q(int levelQ, int levelP, Poly *p1Q, Poly *p1P, Poly *p2Q)
{
    uint64_t N = p1P->Coeffs[0].size();
    uint64_t logN = log(N) / log(2);

#pragma omp parallel for collapse(2)
        for (auto n = 0; n < N; n++) {
            for (auto i = 0; i < levelP + 1; i++) {
            uint64_t *tmpi = baseConvPQ2Q + (i << logN);
            MulModBarrett(tmpi[n], p1P->Coeffs[i][n], pHatInvModp[i], pVec[i], ringP->twoqVec[i], ringP->qrVecHigh[i],
                ringP->qrVecLow[i]);
        }
    }
#pragma omp parallel for collapse(2)
        for (auto n = 0; n < N; n++) {
            for (int i = 0; i < levelQ + 1; i++) {
            uint64_t tt = baseConvPQ2Q[n];
            __uint128_t sum = static_cast<__uint128_t>(tt) * pHatModq[0][i];
            for (auto m = 1; m < levelP + 1; m++) {
                tt = baseConvPQ2Q[n + (m << logN)];
                sum += static_cast<__uint128_t>(tt) * pHatModq[m][i];
            }
            ModBarrett(polypoolQ->Coeffs[i][n], sum, qVec[i], ringQ->twoqVec[i], ringQ->qrVecHigh[i],
                ringQ->qrVecLow[i]);

            int64_t mid = int64_t(p1Q->Coeffs[i][n] - polypoolQ->Coeffs[i][n]);
            while (mid < 0) {
                mid += qVec[i];
            }
            MulModBarrett(p2Q->Coeffs[i][n], mid, PInvModq[i], qVec[i], ringQ->twoqVec[i], ringQ->qrVecHigh[i],
                ringQ->qrVecLow[i]);
        }
    }
}
/*
 * @brief 将基于 P 的多项式转换为基于 Q 的多项式。
 *
 * 该函数实现了将多项式 `p1P`（分别基于 P）转换为一个基于 Q 的多项式 `p2Q`。
 *
 * @param levelQ Q 级数（转换级数）。
 * @param levelP P 级数（转换级数）。
 * @param p1P 输入多项式 p1P，基于 P。
 * @param p2Q 输出多项式 p2Q，转换后的结果，基于 Q。
 */
void FastBasisExtender::BaseConvertP2Q(int levelQ, int levelP, Poly *p1P, Poly *p2Q)
{
    uint64_t N = p1P->Coeffs[0].size();
#pragma omp parallel for
    for (int n = 0; n < N; n++) {
        uint64_t ra[BASECONVERT];
        float vv = 0;
        for (auto i = 0; i < levelP + 1; i++) {
            MulModShoup(ra[i], p1P->Coeffs[i][n], pHatInvModp[i], pHatInvModp_shoup[i], pVec[i]);
            vv += ra[i] / static_cast<float>(pVec[i]);
        }

        for (int i = 0; i < levelQ + 1; i++) {
            __uint128_t sum = 0;
            for (auto m = 0; m < levelP + 1; m++) {
                __uint128_t tt = ra[m];
                sum += tt * pHatModq[m][i];
            }
            uint64_t PQMul = static_cast<int>(vv) * PModq[i];
            ModBarrett(p2Q->Coeffs[i][n], sum - PQMul, qVec[i], ringQ->twoqVec[i], ringQ->qrVecHigh[i],
                ringQ->qrVecLow[i]);
        }
    }
}

